from model.build_BiSeNet import BiSeNet
import torch
import glob
import sys
sys.path.insert(0, '.')


torch.set_grad_enabled(False)


def tf(INPUT_PATH="/project/train/models/best_dice_loss.pth"):

    net = BiSeNet(2, "resnet18")
    net.load_state_dict(torch.load(
        INPUT_PATH, map_location='cpu'), strict=False)
    net.eval()
    # net.module.load_state_dict(torch.load("a"))

    dummy_input = torch.randn(1, 3, 720, 960)
    input_names = ['input']
    output_names = ['output']
    torch.onnx.export(net, dummy_input, INPUT_PATH[-3]+"onnx",
                      input_names=input_names, output_names=output_names,
                      verbose=False, opset_version=11)


pths = glob.glob("/project/train/models/*.pth")
for p in pths:
        print(p)
        tf(p)
